package com.github.cyf.waf;

import com.github.cyf.common.model.Rs;
import com.github.cyf.common.util.JacksonUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.tomcat.util.http.fileupload.servlet.ServletFileUpload;
import org.springframework.util.CollectionUtils;
import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;

/**
 * web 防火墙
 * @author chenyifan
 * @create 2024-11-13 10:36
 */
@Slf4j
@Data
public class WafFilter implements Filter {

    /**
     * 排除链接
     */
    private List<String> excludes = new ArrayList<>();

    private MultipartRequestChecker multipartRequestChecker = new MultipartRequestChecker();

    private List<Pattern> excludePatterns;

    public void setExcludes(List<String> excludes) {
        this.excludes = excludes;
    }

    @Override
    public void init(FilterConfig filterConfig) {
        excludePatterns = new ArrayList<>();
        if (!CollectionUtils.isEmpty(excludes)) {
            for (String pattern : excludes) {
                // ^：代表字符串的开始，确保模式匹配是从目标字符串的最开始位置进行的。
                excludePatterns.add(Pattern.compile("^" + pattern));
            }
        }
        log.info("WafFilter init");
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;

        // 设置安全响应头
        setSecurityHeaders(response);

        // 处理文件上传请求，脚本攻击检测
        if (ServletFileUpload.isMultipartContent(request)) {
            Rs checkResult = multipartRequestChecker.check(request);
            if (Rs.CODE_SUCCESS.equals(checkResult.getCode())) {
                log.warn(" WafFilter exception, multipart requestURL: " + request.getRequestURL());
                response.setStatus(HttpServletResponse.SC_OK);
                // 防止json 中文乱码
                response.setContentType("application/json; charset=UTF-8");
                PrintWriter writer = response.getWriter();
                writer.write(JacksonUtil.toJsonString(checkResult));
                writer.flush();
            }
            chain.doFilter(servletRequest, servletResponse);
            return;
        }

        // 拦截请求，脚本攻击转义
        if (shouldHandleRequest(request)) {
            try {
                //Request请求过滤
                chain.doFilter(new WafRequestWrapper(request), servletResponse);
                return;
            } catch (Exception e) {
                log.error(" WafFilter exception , requestURL: " + request.getRequestURL(), e);
            }
        }

        chain.doFilter(servletRequest, servletResponse);
    }

    @Override
    public void destroy() {
        log.warn(" WafFilter destroy .");
    }

    private void setSecurityHeaders(HttpServletResponse response) {
        // 是否可以在iframe显示视图： DENY=不可以 | SAMEORIGIN=同域下可以 | ALLOW-FROM uri=指定域名下可以
        response.setHeader("X-Frame-Options", "SAMEORIGIN");
        // 禁用浏览器内容嗅探
        response.setHeader("X-Content-Type-Options", "nosniff");
        // 是否启用浏览器默认XSS防护： 0=禁用 | 1=启用 | 1; mode=block 启用, 并在检查到XSS攻击时，停止渲染页面
        response.setHeader("X-XSS-Protection", "1; mode=block");
    }

    private boolean shouldHandleRequest(HttpServletRequest request) {
        if (CollectionUtils.isEmpty(excludes)) {
            return true;
        }
        String url = request.getServletPath();
        for (Pattern pattern : excludePatterns) {
            if (pattern.matcher(url).find()) {
                return false;
            }
        }
        return true;
    }
}